# encoding: utf-8

import torch
from transformers import pipeline
from transformers import BertTokenizer
from transformers import BertForTokenClassification


model_name = "ner_roberta"

device = torch.device("cuda:0") if torch.cuda.is_available() else "cpu"

net = BertForTokenClassification.from_pretrained(
    model_name
)
net.to(device)

tokenizer = BertTokenizer.from_pretrained(model_name)

recognizer = pipeline("token-classification",
                      model=net, tokenizer=tokenizer, aggregation_strategy='simple')

text = "【全网低价】得力(deli)S01中性笔签字笔 0.5mm子弹头经典办公按动笔水笔 黑色 12支/盒"

result1 = recognizer(text)
print("result1=", result1)

result2 = recognizer([text, text])
print("result2=", result2)

"""
result1= [{'entity_group': 'BRAND', 'score': 0.7180295, 'word': '得 力 ( deli', 'start': None, 'end': None}]
result2= [[{'entity_group': 'BRAND', 'score': 0.7180295, 'word': '得 力 ( deli', 'start': None, 'end': None}], [{'entity_group': 'BRAND', 'score': 0.7180295, 'word': '得 力 ( deli', 'start': None, 'end': None}]]


"""